import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
import time
import skimage.metrics as Metrics


FILE_PREFIX = ""

def get_time_stamp():
    return str(time.strftime("%Y-%m-%d-%H%M%S", time.localtime()))

def luckyprint(data, p=1):
    seed = np.random.random()
    if seed < p:
        print(data)

def simu_simple_data(f=20, h=100, w=200):
    # f, w, h = 10, 200, 100
    data = np.zeros([f, h, w])
    for frame_index in range(f):
        data_1d = [0.45*(np.sin(5*np.pi*x/w+2*np.pi*frame_index/f) + 1) for x in range(w)]
        data_2d = np.array([data_1d for y in range(h)])
        data[frame_index] = data_2d
    return data

def simu_random_data(f=20, h=100, w=200):
    # f, w, h = 10, 200, 100
    data = np.zeros([f, h, w])
    for frame_index in range(f):
        data_2d = (np.random.rand(h, w)*0.1).tolist()
        data[frame_index] = data_2d
    return data

def show_video(data):
    plt.cla()
    for frame_idx in range(data.shape[1]):
        plt.imshow(data[0][frame_idx])
        plt.pause(0.05)
    plt.cla()
    

def read_tiff_image(path, type=".tiff", crop=None, size=None, BGR=True):
    # 8 bit image
    # luckyprint(path, p=0.02)
    if (os.path.splitext(path)[-1] == type):
        img = cv2.imread(path, 1) 
        if crop:
            img = img[crop[0]:crop[1],crop[2]:crop[3],:]
        if size != None:
            img = cv2.resize(img, (size[1], size[0]))
        if BGR:
            img = img[:, :, (2, 1, 0)]
        
        return img.astype(np.float32)
    else:
        return np.array([[0]])
    


def get_all_image(path, start_idx=0, max_num=100000, imgsize=(0, 0)):
    list_dir = sorted(os.listdir(path))
    list_dir = list_dir[start_idx:min(start_idx+max_num, len(list_dir))]
    img_list = [read_tiff_image(path+img_path, size=imgsize) for img_path in list_dir]
    return np.array(img_list)

def get_low_light_data():
    pass

# [images] should be 2-d list of numpy images
def show_bright_images(images, channel_first=True, figwidth=40, 
                       filename_prefix=FILE_PREFIX, enhanced=True, 
                       BGR=False, closeFig=True):
    if channel_first:
        figheight = figwidth*len(images)/len(images[0])*images[0][0].shape[1]/images[0][0].shape[2]
    else:
        figheight = figwidth*len(images)/len(images[0])*images[0][0].shape[0]/images[0][0].shape[1]
    
    
    
    fig = plt.figure(figsize=(figwidth, figheight))
    for row_idx in range(len(images)):
        for col_idx in range(len(images[row_idx])):
            image = images[row_idx][col_idx]
            fig.tight_layout()
            plt.subplots_adjust(wspace=0, hspace=0)
            plt.subplot(len(images), len(images[0]), row_idx*len(images[0])+col_idx+1)
            if np.mean(image) != 0 and enhanced:
                image = image/np.mean(image)*0.5
            if channel_first and len(image.shape)==3: # only color image has 3 channels
                image = image.transpose(1, 2, 0)
                
            if BGR:
                b,g,r = cv2.split(image)
                image = cv2.merge([r,g,b])
            if len(image.shape)==2:
                plt.imshow(np.clip(image, 0, 1), cmap='gray')
            else:
                plt.imshow(np.clip(image, 0, 1))
                
            plt.axis('off')
    
    plt.savefig("./pic/"+filename_prefix+get_time_stamp()+".png")
    if closeFig:
        plt.close()
    # plt.show()
    
def show_bright_image(image,channel_first=True):
    image = image/np.mean(image)*0.5
    plt.imshow(np.clip(image.transpose(1, 2, 0) if channel_first else image, 0, 1))
    plt.show()

def img_resize(img, size):
    return cv2.resize(img, size, interpolation=cv2.INTER_CUBIC)

def compare_video(vid, gtr, channelfirst=True):
    frame_num = vid.shape[0]
    psnr = []
    ssim = []
    if channelfirst:
        vid = np.transpose(vid, (0, 2, 3, 1))
        gtr = np.transpose(gtr, (0, 2, 3, 1))
    
    for i in range(frame_num):
        psnr.append(Metrics.peak_signal_noise_ratio(gtr[i], vid[i]))
        ssim.append(Metrics.structural_similarity(gtr[i], vid[i], multichannel=True))

    return np.mean(psnr), np.mean(ssim)
